#=
测试对phi求导
=#



include("../src/qmcmary.jl")
using ..qmcmary

using Test
using ReverseDiff: jacobian, jacobian!
using LinearAlgebra

BLAS.set_num_threads(1)

function sgn_scratch(ss::ScrollSVD{T}) where T
    siz = size(ss.B[end])
    ptr = findfirst(ss.L)
    if isnothing(ptr)
        VL = Diagonal(ones(siz[1]))
        DL = VL
        UL = VL
        UR, DR, VR = ss.F[end].U, Diagonal(ss.F[end].S), ss.F[end].Vt
    elseif ptr == 1
        VL, DL, UL = ss.F[1].U, Diagonal(ss.F[1].S), ss.F[1].Vt
        UR = Diagonal(ones(siz[1]))
        DR = UR
        VR = UR
        ##上面这个数值稳定性会突然出问题
        ##用下面这个
        ##VL DL UL UR=I DR=I VR=I -> I I I UR=VL DR=DL VR=UL
        #VL = Diagonal(ones(siz[1]))
        #DL = VL
        #UL = VL
        #UR, DR, VR = ss.F[1].U, Diagonal(ss.F[1].S), ss.F[1].Vt
    else
        VL, DL, UL = ss.F[ptr].U, Diagonal(ss.F[ptr].S), ss.F[ptr].Vt
        UR, DR, VR = ss.F[ptr-1].U, Diagonal(ss.F[ptr-1].S), ss.F[ptr-1].Vt
    end
    #gtt = inv(Diagonal(ones(siz[1]))+UR*DR*VR*VL*DL*UL)
    #M = inv(UL*UR) + DR*(VR*VL)*DL
    #Fm = svd(M)
    #gtt = inv(Fm.Vt*UL)*inv(Diagonal(Fm.S))*inv(UR*Fm.U)
    #
    DLS = Diagonal(ones(Float64, siz[1]))
    DLB = Diagonal(ones(Float64, siz[1]))
    DRS = Diagonal(ones(Float64, siz[1]))
    DRB = Diagonal(ones(Float64, siz[1]))
    for i in Base.OneTo(siz[1])
        if DL[i, i] > 1.0
            DLB[i, i] = DL[i, i]
            DLS[i, i] = 1.0
        else
            DLS[i, i] = DL[i, i]
            DLB[i, i] = 1.0
        end
        if DR[i, i] > 1.0
            DRB[i, i] = DR[i, i]
            DRS[i, i] = 1.0
        else
            DRS[i, i] = DR[i, i]
            DRB[i, i] = 1.0
        end
    end
    #
    M = inv(DRB)*adjoint(UL*UR)*inv(DLB) + DRS*(VR*VL)*DLS
    Fm = svd(M, alg=LinearAlgebra.QRIteration())
    #ML = adjoint(UL)*inv(DLB)*adjoint(Fm.Vt)
    #MR = adjoint(Fm.U)*inv(DRB)*adjoint(UR)
    #gtt = ML*inv(Diagonal(Fm.S))*MR
    #增加计算phase
    sgn = det(Fm.Vt*UL)*det(UR*Fm.U)
    sgn = log(abs(sgn))
    for sval in Fm.S
        sgn += log(sval)
    end
    for sval in diag(DLB)
        sgn += log(sval)
    end
    for sval in diag(DRB)
        sgn += log(sval)
    end
    return sgn
end


"""
b矩阵数值求导
"""
function pbpϕ_example()
    L = 4
    Nx = L^2
    dt = 0.1
    #
    tp = 0.25
    barehk = lattice_tprim_square(ComplexF64, L, -1.0+0.0im, tp+0.0im)
    #
    Ui = 6*ones(Nx)
    phis = rand(Nx)
    #4分量的hs变换
    #这里只处理这一种情况
    #a = -Ui/2
    #A = (n↑ - n↓ + ϕ)^2
    #最后会出现额外的-Uϕn↑ + Uϕn↓，需要补充额外的化学势Uϕn↑ - Uϕn↓
    # H = (-t∑) + (-Ui/2)A + Uϕn↑ - Uϕn↓
    hk = kron([1 0; 0 1], barehk)
    pm = kron([1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    Nt = 50
    Ng = 5
    hscfg = Matrix{Int}(undef, Nt, Nx)
    for bi in Base.OneTo(Nt)
        cfg = floor.(4rand(Nx)) .- 2
        cfg = map(x -> x >= 0 ? x+1 : x, cfg)
        hscfg[bi, :] .= cfg
    end
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats1, ss1 = initialize_SS_Quad1(Nt, Ng, sp, hscfg, phis)
    logS1 = sgn_scratch(ss1)
    #
    ni = 1
    phis[ni] += 0.001
    hk = kron([1 0; 0 1], barehk)
    pm = kron([1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats2, ss2 = initialize_SS_Quad1(Nt, Ng, sp, hscfg, phis)
    logS2 = sgn_scratch(ss2)
    #
    #
    tapemap = pbpϕ_calc_tapemap(sp, phis)
    adr, adi = pbpϕ_calc_xi(sp, hscfg[:, ni], ni, phis; tapemap=tapemap)
    #
    println(allbmats2[1][[ni, ni+Nx], :] - allbmats1[1][[ni, ni+Nx], :])
    println(hscfg[1, 1], adr[1, :, :])
    #
    println(allbmats2[5][[ni, ni+Nx], :] - allbmats1[5][[ni, ni+Nx], :])
    println(hscfg[5, 1], adr[5, :, :])
    #
    println(-logS2 + logS1)
    #
    phibar = meas_gradϕ(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=true, ignorez=true)
    println(phibar[ni])
end


function pwpϕ_example()
    #
    L = 4
    Nx = L^2
    dt = 0.1
    #
    tp = 0.25
    barehk = lattice_tprim_square(ComplexF64, L, -1.0+0.0im, tp+0.0im)
    #
    Ui = 4*ones(Nx)
    phis = rand(Nx)
    #4分量的hs变换
    #这里只处理这一种情况
    #a = -Ui/2
    #A = (n↑ - n↓ + ϕ)^2
    #最后会出现额外的-Uϕn↑ + Uϕn↓，需要补充额外的化学势Uϕn↑ - Uϕn↓
    hk = kron([1 0; 0 1], barehk)
    pm = kron([1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    Nt = 50
    Ng = 5
    hscfg = Matrix{Int}(undef, Nt, Nx)
    for bi in Base.OneTo(Nt)
        cfg = floor.(4rand(Nx)) .- 2
        cfg = map(x -> x >= 0 ? x+1 : x, cfg)
        hscfg[bi, :] .= cfg
    end
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats1, ss1 = initialize_SS_Quad1(Nt, Ng, sp, hscfg, phis)
    logS1 = sgn_scratch(ss1)
    #循环每一个hs，将权重增加到对数上
    ηdict = Dict(
        -2 => -3.301360247771569, 
        -1 => -1.049295246550581,
        1 => 1.049295246550581,
        2 => 3.301360247771569
    )
    γdict = Dict(
        -2 => 0.18350341907227408, 
        -1 => 1.8164965809277258,
        1 => 1.8164965809277258,
        2 => 0.18350341907227408
    )
    for bi in Base.OneTo(Nt); for xi in Base.OneTo(Nx)
        logS1 += sqrt(dt*Ui[xi]/2)*ηdict[hscfg[bi, xi]]*phis[xi]
        logS1 += log(γdict[hscfg[bi, xi]])
    end; end
    #
    ni = 1
    phis[ni] += 0.001
    hk = kron([1 0; 0 1], barehk)
    pm = kron([1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats2, ss2 = initialize_SS_Quad1(Nt, Ng, sp, hscfg, phis)
    logS2 = sgn_scratch(ss2)
    for bi in Base.OneTo(Nt); for xi in Base.OneTo(Nx)
        logS2 += sqrt(dt*Ui[xi]/2)*ηdict[hscfg[bi, xi]]*phis[xi]
        logS2 += log(γdict[hscfg[bi, xi]])
    end; end
    #
    tapemap = pbpϕ_calc_tapemap(sp, phis)
    #adr, adi = pbpϕ_calc_xi(sp, hscfg[:, ni], ni, phis; tapemap=tapemap)
    #
    println(-logS2 + logS1)
    phibar = meas_gradϕ(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=true, ignorez=true)
    println(phibar[ni])
    phibar = meas_gradϕ(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=false, ignorez=true)
    println(phibar[ni])
    #
    dlnZ = 0.0
    dlnZ = Ui[ni]*0.5*Nt*dt*phis[ni]^2 - Ui[ni]*0.5*Nt*dt*(phis[ni]-0.001)^2
    phibar = meas_gradϕ(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=false, ignorez=false)
    println(-logS2 + logS1 + dlnZ)
    println(phibar[ni])
end


"""
b矩阵数值求导
"""
function pbpϕ2_example()
    L = 4
    Nx = L^2
    dt = 0.1
    #
    tp = 0.25
    barehk = lattice_tprim_square(ComplexF64, L, -1.0+0.0im, tp+0.0im)
    #
    Ui = 6*ones(Nx)
    phis = rand(Nx)
    #4分量的hs变换
    #这里只处理这一种情况
    #a = Ui/2
    #A = (n↑ + n↓ - 1 + ϕ)^2
    #最后会出现额外的+Uϕn↑ + Uϕn↓，需要补充额外的化学势-Uϕn↑ - Uϕn↓
    # H = (-t∑) + (Ui/2)A - Uϕn↑ - Uϕn↓
    hk = kron([1 0; 0 1], barehk)
    pm = kron([-1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    Nt = 50
    Ng = 5
    hscfg = Matrix{Int}(undef, Nt, Nx)
    for bi in Base.OneTo(Nt)
        cfg = floor.(4rand(Nx)) .- 2
        cfg = map(x -> x >= 0 ? x+1 : x, cfg)
        hscfg[bi, :] .= cfg
    end
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats1, ss1 = initialize_SS_Quad2(Nt, Ng, sp, hscfg, phis)
    logS1 = sgn_scratch(ss1)
    #
    ni = 1
    phis[ni] += 0.001
    hk = kron([1 0; 0 1], barehk)
    pm = kron([-1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats2, ss2 = initialize_SS_Quad2(Nt, Ng, sp, hscfg, phis)
    logS2 = sgn_scratch(ss2)
    #
    #
    tapemap = pbpϕ2_calc_tapemap(sp, phis)
    adr, adi = pbpϕ2_calc_xi(sp, hscfg[:, ni], ni, phis; tapemap=tapemap)
    #
    println(allbmats2[1][[ni, ni+Nx], :] - allbmats1[1][[ni, ni+Nx], :])
    println(hscfg[1, 1], adr[1, :, :])
    #
    println(allbmats2[5][[ni, ni+Nx], :] - allbmats1[5][[ni, ni+Nx], :])
    println(hscfg[5, 1], adr[5, :, :])
    #
    println(-logS2 + logS1)
    #
    phibar = meas_gradϕ2(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=true)
    println(phibar[ni])
end


function pwpϕ2_example()
    #
    L = 4
    Nx = L^2
    dt = 0.1
    #
    tp = 0.25
    barehk = lattice_tprim_square(ComplexF64, L, -1.0+0.0im, tp+0.0im)
    #
    Ui = 4*ones(Nx)
    phis = rand(Nx)
    #4分量的hs变换
    #这里只处理这一种情况
    #a = -Ui/2
    #A = (n↑ - n↓ + ϕ)^2
    #最后会出现额外的-Uϕn↑ + Uϕn↓，需要补充额外的化学势Uϕn↑ - Uϕn↓
    hk = kron([1 0; 0 1], barehk)
    pm = kron([-1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    Nt = 50
    Ng = 5
    hscfg = Matrix{Int}(undef, Nt, Nx)
    for bi in Base.OneTo(Nt)
        cfg = floor.(4rand(Nx)) .- 2
        cfg = map(x -> x >= 0 ? x+1 : x, cfg)
        hscfg[bi, :] .= cfg
    end
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats1, ss1 = initialize_SS_Quad2(Nt, Ng, sp, hscfg, phis)
    logS1 = sgn_scratch(ss1)
    #循环每一个hs，将权重增加到对数上
    ηdict = Dict(
        -2 => -3.301360247771569, 
        -1 => -1.049295246550581,
        1 => 1.049295246550581,
        2 => 3.301360247771569
    )
    γdict = Dict(
        -2 => 0.18350341907227408, 
        -1 => 1.8164965809277258,
        1 => 1.8164965809277258,
        2 => 0.18350341907227408
    )
    for bi in Base.OneTo(Nt); for xi in Base.OneTo(Nx)
        logS1 += log(abs(exp(sqrt(complex(-dt*Ui[xi]/2, 0.0))*ηdict[hscfg[bi, xi]]*phis[xi])))
        logS1 += log(γdict[hscfg[bi, xi]])
    end; end
    #
    ni = 1
    phis[ni] += 0.001
    hk = kron([1 0; 0 1], barehk)
    pm = kron([-1 0; 0 -1], Ui.*Diagonal(phis))
    hk = hk + pm
    sp = default_splitting(Nt, hk, Ui; Z2=false)
    sslen, bmats, allbmats2, ss2 = initialize_SS_Quad2(Nt, Ng, sp, hscfg, phis)
    logS2 = sgn_scratch(ss2)
    for bi in Base.OneTo(Nt); for xi in Base.OneTo(Nx)
        logS2 += log(abs(exp(sqrt(complex(-dt*Ui[xi]/2, 0.0))*ηdict[hscfg[bi, xi]]*phis[xi])))
        logS2 += log(γdict[hscfg[bi, xi]])
    end; end
    #
    tapemap = pbpϕ2_calc_tapemap(sp, phis)
    #adr, adi = pbpϕ_calc_xi(sp, hscfg[:, ni], ni, phis; tapemap=tapemap)
    #
    println(-logS2 + logS1)
    phibar = meas_gradϕ2(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=true, ignorez=true)
    println(phibar[ni])
    phibar = meas_gradϕ2(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=false, ignorez=true)
    println(phibar[ni])
    #
    dlnZ = 0.0
    dlnZ = -Ui[ni]*0.5*Nt*dt*(phis[ni]-1)^2 +Ui[ni]*0.5*Nt*dt*(phis[ni]-0.001-1)^2
    phibar = meas_gradϕ2(ss1, allbmats1, sp, hscfg, phis; tapemap=tapemap, ignorew=false, ignorez=false)
    println(-logS2 + logS1 + dlnZ)
    println(phibar[ni])
end


#pbpϕ_example()
#pwpϕ_example()
#pbpϕ2_example()
pwpϕ2_example()
